import json
import time
import os, sys
import gradio as gr
import requests
import warnings

warnings.filterwarnings("ignore")

# 定义变量
URL = "http://127.0.0.1:8188/prompt"
INPUT_DIR = "D:\PycharmProjects\ComfyUI\ComfyUI\input"
OUTPUT_DIR = "D:\PycharmProjects\ComfyUI\ComfyUI\output"

# 工具模块路径
UTIL_PATH = r"./utils/module_table.json"

demopath = ""
test_workflow_path= r"/workflows/text_to_img.json"

def open_utils_table(utils_path=UTIL_PATH, mode="tuples"):
    with open(utils_path, "r", encoding="utf-8") as file_json:
        if mode == "array":
            return list(json.load(file_json))
        elif mode == "tuples":
            return json.load(file_json)
        else:
            print("函数参数报错")


# 全局工具模块表
UTILS_TABLE_NAME = open_utils_table(UTIL_PATH, "array")
UTILS_TABLE = open_utils_table(UTIL_PATH)
print("UTILS_TABLE_NAME", UTILS_TABLE_NAME)
print("UTILS_TABLE", UTILS_TABLE)


# 加载json文件解析节点v3
def load_json_to_modules3():
    with gr.Blocks(analytics_enabled=False) as demo:
        with open(test_workflow_path, "r", encoding="utf-8") as file_json:
            inputarr = []
            outputarr = []
            idarr = []
            inputs1 = []
            output1 = []
            modules = json.load(file_json)
        gr.Markdown("<div align='center'>"
                    " <h1>ComfyUI_Easy_Link</h1>"
                    "<h3>作者：Xik</h3>"
                    "</div>")
        with gr.Tab("图生图"):
            gr.Text()

        with gr.Tab("文生图"):
            gr.Text()

        btn = gr.Button("开始绘画")
        btn.click(
            fn=generate_image,
            inputs=inputs1,
            outputs=output1
        )

    return demo


# 加载json文件解析节点v2
def load_json_to_modules2():
    with open(demopath, "r", encoding="utf-8") as file_json:
        inputarr = []
        outputarr = []
        idarr = []
        inputs1 = []
        output1 = []
        modules = json.load(file_json)
        with gr.Blocks(analytics_enabled=False) as demo:
            gr.Markdown("<div align='center'>"
                        " <h1>ComfyUI_Easy_Link</h1>"
                        "<h3>作者：谢凯军</h3>"
                        "</div>")

            for i in list(modules):
                if modules[i]['class_type'] in UTILS_TABLE_NAME:
                    if UTILS_TABLE[modules[i]['class_type']]["gr_input"] != "":
                        inputarr.append(UTILS_TABLE[modules[i]['class_type']]["gr_input"])
                        idarr.append(i)
                        if UTILS_TABLE[modules[i]['class_type']]["gr_input"] == "text":
                            text1 = gr.Text()
                            inputs1.append(text1)
                    if UTILS_TABLE[modules[i]['class_type']]["gr_out"] != "":
                        outputarr.append(UTILS_TABLE[modules[i]['class_type']]["gr_out"])
                        if UTILS_TABLE[modules[i]['class_type']]["gr_out"] == "image":
                            text2 = gr.Image()
                            output1.append(text2)
            btn = gr.Button("开始绘画")

            btn.click(
                fn=generate_image,
                inputs=inputs1,
                outputs=output1
            )

    return demo


# 加载json文件解析节点v1
def load_json_to_modules(json_path):
    inputarr = []
    outputarr = []
    with open(json_path, "r", encoding="utf-8") as file_json:
        modules = json.load(file_json)

        for i in list(modules):
            if modules[i]['class_type'] in UTILS_TABLE_NAME:
                if UTILS_TABLE[modules[i]['class_type']]["gr_input"] != "":
                    inputarr.append(UTILS_TABLE[modules[i]['class_type']]["gr_input"])
                if UTILS_TABLE[modules[i]['class_type']]["gr_out"] != "":
                    outputarr.append(UTILS_TABLE[modules[i]['class_type']]["gr_out"])

    return inputarr, outputarr


# inputarr1,outputarr1= load_json_to_modules(r"D:\PycharmProjects\Comfyui_link\ComfyUI_easy_link\workflows\text_to_img.json")
# print("inputarr1", inputarr1,outputarr1)


# 定义获取最新图像的逻辑方法，用于之后下面函数的调用
def get_latest_image(folder):
    files = os.listdir(folder)
    image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder, x)))
    latest_image = os.path.join(folder, image_files[-1]) if image_files else None
    return latest_image


# 开始获取请求进行编码
def start_queue(prompt_workflow):
    p = {"prompt": prompt_workflow}
    data = json.dumps(p).encode('utf-8')
    requests.post(URL, data=data)


# 开始生成图像，前端UI定义所需变量传递给json
def generate_image(*prompt_text):
    with open(demopath, "r", encoding="utf-8") as file_json:
        prompt = json.load(file_json)
        print(prompt_text)
        index = 0
        for i in list(prompt):
            if prompt[i]['class_type'] in UTILS_TABLE_NAME:
                if UTILS_TABLE[prompt[i]['class_type']]["gr_input"] != "" or UTILS_TABLE[prompt[i]['class_type']][
                    "input_label"] != "":
                    print("deaasdas", UTILS_TABLE[prompt[i]['class_type']]["input_label"])
                    prompt[i]["inputs"][UTILS_TABLE[prompt[i]['class_type']]["input_label"]] = f"{prompt_text[index]}"
                    index += 1

        # prompt["6"]["inputs"]["text"] = f"digital artwork of a {prompt_text}"
        # prompt['9']['inputs']['filename_prefix'] = str(time.time())

    previous_image = get_latest_image(OUTPUT_DIR)  # 推理出的最新输出图像保存到指定的OUTPUT_DIR变量路径

    start_queue(prompt)

    # 这是一个循环获取指定路径的最新图像，休眠·一秒钟后继续循环
    while True:
        latest_image = get_latest_image(OUTPUT_DIR)
        if latest_image != previous_image:
            return latest_image

        time.sleep(1)


if __name__ == "__main__":
    demo = load_json_to_modules3()
    demo.queue()
    #demo.launch(auth=("admin", "pass1234"))
    demo.launch()
